import numpy as np
import warnings
import os

def iou(box1, box2):
    """计算两个边界框的交并比(IoU)"""
    x1, y1, x2, y2 = box1
    x1_, y1_, x2_, y2_ = box2
    # 计算交集坐标
    inter_x1 = max(x1, x1_)
    inter_y1 = max(y1, y1_)
    inter_x2 = min(x2, x2_)
    inter_y2 = min(y2, y2_)
    # 计算交集面积
    inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)
    # 计算两个框的面积
    area1 = (x2 - x1) * (y2 - y1)
    area2 = (x2_ - x1_) * (y2_ - y1_)
    # 计算并集面积
    union_area = area1 + area2 - inter_area
    return inter_area / union_area if union_area > 0 else 0.0

def gen_golden_data_simple():
    warnings.filterwarnings('ignore', category=RuntimeWarning)
    
    # 创建输入/输出目录
    os.makedirs('input', exist_ok=True)
    os.makedirs('output', exist_ok=True)
    
    # 1. 生成输入数据（boxes和scores）
    all_boxes = np.array([
        [10.0, 10.0, 20.0, 20.0],   # 0
        [12.0, 12.0, 18.0, 18.0],   # 1
        [15.0, 15.0, 19.0, 19.0],   # 2
        [14.0, 14.0, 18.0, 18.0],   # 3
        [18.0, 18.0, 25.0, 25.0],   # 4
        [25.0, 10.0, 35.0, 20.0],   # 5
        [10.0, 25.0, 20.0, 35.0],   # 6
        [20.0, 20.0, 23.0, 23.0],   # 7
        [21.0, 21.0, 24.0, 24.0],   # 8
        [30.0, 30.0, 40.0, 40.0],   # 9
        [45.0, 10.0, 55.0, 20.0],   # 10
        [10.0, 40.0, 20.0, 50.0],   # 11
        [60.0, 30.0, 70.0, 40.0],   # 12
        [30.0, 50.0, 40.0, 60.0],   # 13
        [75.0, 10.0, 85.0, 20.0],   # 14
        [10.0, 65.0, 20.0, 75.0],   # 15
        [50.0, 60.0, 60.0, 70.0],   # 16
        [80.0, 40.0, 90.0, 50.0],   # 17
        [30.0, 75.0, 40.0, 85.0],   # 18
        [95.0, 20.0, 105.0, 30.0]   # 19
    ], dtype=np.float32)
    
    all_scores = np.array([
        0.99, 0.98, 0.97, 0.96, 0.95,
        0.94, 0.93, 0.92, 0.91, 0.90,
        0.89, 0.88, 0.87, 0.86, 0.85,
        0.84, 0.83, 0.82, 0.81, 0.80
    ], dtype=np.float32)
    
    # 写入输入文件
    all_boxes.tofile('input/input_x.bin')
    all_scores.tofile('input/input_y.bin')
    
    # 2. 生成黄金标准数据（NMS预期输出）
    iou_threshold = 0.1
    scores_threshold = 0.1
    keep_indices = []
    golden = np.zeros(len(all_boxes), dtype=np.uint8)
    
    for i in range(len(all_boxes)):
        # 过滤低于分数阈值的框
        if all_scores[i] < scores_threshold:
            golden[i] = 0
            continue
        
        current_box = all_boxes[i]
        valid = True
        # 与已保留的框计算IoU
        for idx in keep_indices:
            if iou(current_box, all_boxes[idx]) > iou_threshold:
                valid = False
                break
        
        if valid:
            golden[i] = 1
            keep_indices.append(i)
        else:
            golden[i] = 0
    
    # 写入黄金标准文件
    golden.tofile('output/golden.bin')
    print("INFO: 数据生成完成（输入和黄金标准）")

if __name__ == "__main__":
    gen_golden_data_simple()